import json
import os
import numpy as np
import open3d as o3d
from plyfile import PlyData, PlyElement

def read_ply(ply_path):
    # 读取header部分（使用二进制模式解决编码问题）
    header = []
    with open(ply_path, 'rb') as f:  # 修改为二进制模式
        for line_bytes in f:
            line = line_bytes.decode('utf-8').strip()  # 显式解码
            header.append(line)
            if line == 'end_header':
                break
    
    # 读取顶点数据（使用plyfile库原生方法）
    ply_data = PlyData.read(ply_path)
    vertex_element = ply_data['vertex']
    vertex_data = []
    
    # 转换每个顶点的属性为字符串列表
    for v in vertex_element:
        row = []
        for prop in vertex_element.properties:
            val = v[prop.name]
            # 根据数据类型格式化字符串
            if 'float' in prop.val_dtype:
                formatted = f"{val:.6f}".rstrip('0').rstrip('.') if '.' in f"{val:.6f}" else f"{val}"
            elif 'int' in prop.val_dtype or 'uint' in prop.val_dtype:
                formatted = f"{val}"
            else:
                formatted = str(val)
            row.append(formatted)
        vertex_data.append(row)
    
    return header, vertex_data

def write_ply(output_path, header, vertices):
    # 更新顶点数量
    for i, line in enumerate(header):
        if line.startswith('element vertex'):
            header[i] = f'element vertex {len(vertices)}'
            break
    
    with open(output_path, 'w', encoding='utf-8') as f:  # 显式指定编码
        f.write('\n'.join(header) + '\n')
        for v in vertices:
            f.write(' '.join(v) + '\n')

def merge_point_clouds(metadata_path, output_path):
    with open(metadata_path, 'r') as f:
        metadata = json.load(f)
    
    all_vertices = []
    common_header = None
    
    for group in metadata["groups"]:
        obj_path = group["output_file"]
        ply_path = os.path.splitext(obj_path)[0] + ".ply"
        translate = group["transform"]["translate"]
        tx, ty, tz = translate
        
        if not os.path.exists(obj_path) or not os.path.exists(ply_path):
            print(f"Warning: {obj_path} or {ply_path} not found. Skipping.")
            continue
        
        # 读取并采样OBJ模型
        try:
            # 使用Open3D读取OBJ文件
            mesh = o3d.io.read_triangle_mesh(obj_path)
        except Exception as e:
            print(f"Error reading {obj_path}: {str(e)}")
            continue
        
        # 验证网格数据
        if not mesh.has_vertices():
            print(f"Warning: {obj_path} has no vertices. Skipping.")
            continue
        if not mesh.triangles:
            print(f"Warning: {obj_path} has no surface geometry. Skipping.")
            continue
        
        # 均匀采样表面点云
        try:
            target_pcd = mesh.sample_points_uniformly(number_of_points=20000)
        except Exception as e:
            print(f"Sampling failed for {obj_path}: {str(e)}")
            continue
        
        # 读取PLY数据
        try:
            header, vertex_data = read_ply(ply_path)
        except Exception as e:
            print(f"Error reading {ply_path}: {str(e)}")
            continue
        
        if not vertex_data:
            print(f"Warning: {ply_path} has no vertices. Skipping.")
            continue
        
        # 提取有效顶点
        source_points = []
        valid_vertex_indices = []
        for i, v in enumerate(vertex_data):
            if len(v) >= 3:
                try:
                    x = float(v[0])
                    y = float(v[1])
                    z = float(v[2])
                    source_points.append([x, y, z])
                    valid_vertex_indices.append(i)
                except ValueError:
                    continue
        
        if not source_points:
            print(f"Warning: {ply_path} has no valid vertices. Skipping.")
            continue
        
        # 构建源点云
        source_pcd = o3d.geometry.PointCloud()
        source_pcd.points = o3d.utility.Vector3dVector(np.array(source_points))
        
        # ICP配准
        threshold = 1.0
        trans_init = np.eye(4)
        reg_result = o3d.pipelines.registration.registration_icp(
            source_pcd, target_pcd, threshold, trans_init,
            o3d.pipelines.registration.TransformationEstimationPointToPoint(),
            o3d.pipelines.registration.ICPConvergenceCriteria(max_iteration=400))
        
        # 应用变换
        T = reg_result.transformation
        transformed_vertices = []
        for i in valid_vertex_indices:
            v = vertex_data[i]
            try:
                x = float(v[0])
                y = float(v[1])
                z = float(v[2])
            except ValueError:
                continue
            
            # 应用ICP变换
            p_homog = np.array([x, y, z, 1.0])
            p_transformed = T @ p_homog
            
            # 应用全局平移
            x_t = p_transformed[0] + tx
            y_t = p_transformed[1] + ty
            z_t = p_transformed[2] + tz
            
            # 保留其他属性
            new_v = [f"{x_t:.6f}", f"{y_t:.6f}", f"{z_t:.6f}"] + v[3:]
            transformed_vertices.append(new_v)
        
        # 更新公共header
        if not common_header:
            common_header = header.copy()
            # 初始化顶点计数为0，后续统一更新
            for idx, line in enumerate(common_header):
                if line.startswith("element vertex"):
                    common_header[idx] = "element vertex 0"
                    break
        
        all_vertices.extend(transformed_vertices)
    
    if not common_header or len(all_vertices) == 0:
        print("Error: No valid data processed.")
        return
    
    # 更新总顶点数
    for idx, line in enumerate(common_header):
        if line.startswith("element vertex"):
            common_header[idx] = f"element vertex {len(all_vertices)}"
            break
    
    # 写入最终文件
    write_ply(output_path, common_header, all_vertices)
    print(f"Merged point cloud saved to {output_path}")

if __name__ == "__main__":
    merge_point_clouds("output_metadata.json", "merged.ply")